import glob
import json
import ntpath
import random

from param_parser import parameter_parser
from tqdm import tqdm, trange


class DataHandle(object):
    def __init__(self, args):
        self.args = args

    def formatGrahpInfo(self):
        self.train_graphs = glob.glob(self.args.training_graphs + "*.json")
        for trainFilePath in tqdm(self.train_graphs):
            fileName = ntpath.basename(trainFilePath)
            newFilePath = self.args.new_training_graphs + fileName
            newGraphs = []
            with open(trainFilePath) as file:
                for line in file:
                    graph = json.loads(line)
                    newgraph = {'src': graph['src'], 'fname': graph['fname'], 'features': graph['features'],
                                'edges': []}
                    for index, nextPoints in enumerate(graph['succs']):
                        for nextPoint in nextPoints:
                            newgraph['edges'] += [[index, nextPoint]]
                    newGraphs += [newgraph]
            with open(newFilePath, 'a+') as f:
                for graph in newGraphs:
                    json_graph = json.dumps(graph)
                    f.write(json_graph + '\n')

    def generatorGraphPair(self):
        self.new_train_graphs = glob.glob(self.args.new_training_graphs + "*.json")
        file1 = self.new_train_graphs[0]
        file2Index = random.randint(1, len(self.new_train_graphs) - 1)
        file2 = self.new_train_graphs[1]
        graphs1 = []
        graphs2 = []
        with open(file1) as f:
            for line in f:
                graph = json.loads(line)
                newgraph = {'src': graph['src'], 'fname': graph['fname'], 'features': graph['features'],
                            'edges': graph['edges']}
                graphs1 += [newgraph]

        with open(file2) as f:
            for line in f:
                graph = json.loads(line)
                newgraph = {'src': graph['src'], 'fname': graph['fname'], 'features': graph['features'],
                            'edges': graph['edges']}
                graphs2 += [newgraph]
        # print(graphs1)

        graph_pairs = []

        for graph1 in graphs1:
            for graph2 in graphs2:
                if graph1['fname'] == graph2['fname']:
                    graph_pair = {}
                    graph_pair['src1'] = graph1['src']
                    graph_pair['src2'] = graph2['src']
                    graph_pair['fname1'] = graph1['fname']
                    graph_pair['fname2'] = graph2['fname']
                    graph_pair['features1'] = graph1['features']
                    graph_pair['features2'] = graph2['features']
                    graph_pair['edges1'] = graph1['edges']
                    graph_pair['edges2'] = graph2['edges']
                    graph_pair['label'] = 1
                    graph_pairs += [graph_pair]
                    break

        for i in range(len(graphs1)):
            # for i in range(50):
            graph_pair = {}
            graph1 = graphs1[i]
            graph2Index = random.randint(0, len(graphs2) - 1)
            while i == graph2Index:
                graph2Index = random.randint(0, len(graphs2) - 1)
            graph2 = graphs2[graph2Index]
            graph_pair['src1'] = graph1['src']
            graph_pair['src2'] = graph2['src']
            graph_pair['fname1'] = graph1['fname']
            graph_pair['fname2'] = graph2['fname']
            graph_pair['features1'] = graph1['features']
            graph_pair['features2'] = graph2['features']
            graph_pair['edges1'] = graph1['edges']
            graph_pair['edges2'] = graph2['edges']
            graph_pair['label'] = 0
            graph_pairs += [graph_pair]

        train_data_path = self.args.train_data + "train.json"
        with open(train_data_path, "a+") as f:
            for graphPair in graph_pairs:
                if graphPair['edges1'] != [] and graphPair['edges2'] != []:
                    jsonGraph = json.dumps(graphPair)
                    f.write(jsonGraph + '\n')


    def generatorGraphPairTest(self):
        self.new_train_graphs = glob.glob(self.args.new_training_graphs + "*.json")
        file1 = self.new_train_graphs[4]
        file2Index = random.randint(1, len(self.new_train_graphs) - 1)
        file2 = self.new_train_graphs[5]
        graphs1 = []
        graphs2 = []
        with open(file1) as f:
            for line in f:
                graph = json.loads(line)
                newgraph = {'src': graph['src'], 'fname': graph['fname'], 'features': graph['features'],
                            'edges': graph['edges']}
                graphs1 += [newgraph]

        with open(file2) as f:
            for line in f:
                graph = json.loads(line)
                newgraph = {'src': graph['src'], 'fname': graph['fname'], 'features': graph['features'],
                            'edges': graph['edges']}
                graphs2 += [newgraph]
        # print(graphs1)

        graph_pairs = []

        for graph1 in graphs1:
            for graph2 in graphs2:
                if graph1['fname'] == graph2['fname']:
                    graph_pair = {}
                    graph_pair['src1'] = graph1['src']
                    graph_pair['src2'] = graph2['src']
                    graph_pair['fname1'] = graph1['fname']
                    graph_pair['fname2'] = graph2['fname']
                    graph_pair['features1'] = graph1['features']
                    graph_pair['features2'] = graph2['features']
                    graph_pair['edges1'] = graph1['edges']
                    graph_pair['edges2'] = graph2['edges']
                    graph_pair['label'] = 1
                    graph_pairs += [graph_pair]
                    break

        for i in range(len(graphs1)):
            # for i in range(50):
            graph_pair = {}
            graph1 = graphs1[i]
            graph2Index = random.randint(0, len(graphs2) - 1)
            while i == graph2Index:
                graph2Index = random.randint(0, len(graphs2) - 1)
            graph2 = graphs2[graph2Index]
            graph_pair['src1'] = graph1['src']
            graph_pair['src2'] = graph2['src']
            graph_pair['fname1'] = graph1['fname']
            graph_pair['fname2'] = graph2['fname']
            graph_pair['features1'] = graph1['features']
            graph_pair['features2'] = graph2['features']
            graph_pair['edges1'] = graph1['edges']
            graph_pair['edges2'] = graph2['edges']
            graph_pair['label'] = 0
            graph_pairs += [graph_pair]

        train_data_path = self.args.train_data + "test.json"
        with open(train_data_path, "a+") as f:
            for graphPair in graph_pairs:
                if graphPair['edges1'] != [] and graphPair['edges2'] != []:
                    jsonGraph = json.dumps(graphPair)
                    f.write(jsonGraph + '\n')
